package type_alias

import "slices"

// Set 是一个泛型集合，内部使用 map 存储元素。
// 通过结构体封装，便于扩展更多功能。
type Set[T comparable] struct {
	m map[T]struct{}
}

// NewSet 创建并返回一个新的集合。
// capacity 指定 map 的初始容量，传 0 表示默认。
// 可选 elems 参数用于初始化若干元素。
func NewSet[T comparable](capacity int, elems ...T) *Set[T] {
	set := &Set[T]{
		m: make(map[T]struct{}, capacity),
	}
	for _, e := range elems {
		set.m[e] = struct{}{}
	}
	return set
}

// Add 将元素加入集合。
func (s *Set[T]) Add(elem T) {
	s.m[elem] = struct{}{}
}

// AddAll 批量加入元素。
func (s *Set[T]) AddAll(elems ...T) {
	for _, elem := range elems {
		s.m[elem] = struct{}{}
	}
}

// Remove 从集合中移除元素。
func (s *Set[T]) Remove(elem T) {
	delete(s.m, elem)
}

// Exists 判断元素是否存在。
func (s *Set[T]) Exists(elem T) bool {
	_, ok := s.m[elem]
	return ok
}

// Len 返回集合大小。
func (s *Set[T]) Len() int {
	return len(s.m)
}

// Clear 清空集合。
func (s *Set[T]) Clear() {
	clear(s.m)
}

// Range 遍历集合。
// 回调返回 false 时提前终止遍历。
func (s *Set[T]) Range(f func(elem T) bool) {
	for elem := range s.m {
		if !f(elem) {
			return
		}
	}
}

// Elements 返回集合元素的切片副本（无特定顺序）。
func (s *Set[T]) Elements() []T {
	res := make([]T, 0, len(s.m))
	for elem := range s.m {
		res = append(res, elem)
	}
	return res
}

// SortedElements 返回排序后的切片副本，需要调用者提供 less 函数。
func (s *Set[T]) SortedElements(less func(a, b T) bool) []T {
	res := s.Elements()
	slices.SortFunc(res, func(a, b T) int {
		switch {
		case less(a, b):
			return -1
		case less(b, a):
			return 1
		default:
			return 0
		}
	})
	return res
}

// Clone 返回集合的浅拷贝。
func (s *Set[T]) Clone() *Set[T] {
	copySet := NewSet[T](len(s.m))
	for elem := range s.m {
		copySet.m[elem] = struct{}{}
	}
	return copySet
}

// Map 返回底层 map 的浅拷贝。
func (s *Set[T]) Map() map[T]struct{} {
	cp := make(map[T]struct{}, len(s.m))
	for k := range s.m {
		cp[k] = struct{}{}
	}
	return cp
}

// Union 并集：返回包含 a、b 所有元素的新集合。
func Union[T comparable](a, b *Set[T]) *Set[T] {
	res := a.Clone()
	for elem := range b.m {
		res.m[elem] = struct{}{}
	}
	return res
}

// Intersect 交集：返回仅包含 a、b 共同元素的新集合。
func Intersect[T comparable](a, b *Set[T]) *Set[T] {
	var small, large *Set[T]
	if a.Len() <= b.Len() {
		small, large = a, b
	} else {
		small, large = b, a
	}

	res := NewSet[T](small.Len())
	for elem := range small.m {
		if large.Exists(elem) {
			res.m[elem] = struct{}{}
		}
	}
	return res
}

// Difference 差集：返回存在于 a 但不存在于 b 的新集合。
func Difference[T comparable](a, b *Set[T]) *Set[T] {
	res := NewSet[T](a.Len())
	for elem := range a.m {
		if !b.Exists(elem) {
			res.m[elem] = struct{}{}
		}
	}
	return res
}

// SymmetricDifference 对称差：返回只存在于其中一个集合的元素。
func SymmetricDifference[T comparable](a, b *Set[T]) *Set[T] {
	res := NewSet[T](a.Len() + b.Len())
	for elem := range a.m {
		if !b.Exists(elem) {
			res.m[elem] = struct{}{}
		}
	}
	for elem := range b.m {
		if !a.Exists(elem) {
			res.m[elem] = struct{}{}
		}
	}
	return res
}

// IsSubset 判断 a 是否为 b 的子集。
func IsSubset[T comparable](a, b *Set[T]) bool {
	if a.Len() > b.Len() {
		return false
	}
	for elem := range a.m {
		if !b.Exists(elem) {
			return false
		}
	}
	return true
}
